(*

   Copyright 2025, HE Weikun

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

*)

type lietype = E6 | E7 | E8 | F4 | G2

let subscript = function
  | 0 -> "₁"
  | 1 -> "₂"
  | 2 -> "₃"
  | 3 -> "₄"
  | 4 -> "₅"
  | 5 -> "₆"
  | 6 -> "₇"
  | 7 -> "₈"
  | 8 -> "₉"
  | n -> Int.to_string n

let name = function
  | E6 -> "E₆"
  | E7 -> "E₇"
  | E8 -> "E₈"
  | F4 -> "F₄"
  | G2 -> "G₂"

let of_string = function
  | "E6" -> E6
  | "E7" -> E7
  | "E8" -> E8
  | "F4" -> F4
  | "G2" -> G2
  | _ -> raise (Invalid_argument "Dynkin.of_string")

let rank = function
  | E6 -> 6
  | E7 -> 7
  | E8 -> 8
  | F4 -> 4
  | G2 -> 2

let cartan_matrix = function
  | E6 -> [|
      [| 2;  0; -1;  0;  0;  0 |];
      [| 0;  2;  0; -1;  0;  0 |];
      [|-1;  0;  2; -1;  0;  0 |];
      [| 0; -1; -1;  2; -1;  0 |];
      [| 0;  0;  0; -1;  2; -1 |];
      [| 0;  0;  0;  0; -1;  2 |]
    |]
  | E7 -> [|
      [| 2;  0; -1;  0;  0;  0;  0 |];
      [| 0;  2;  0; -1;  0;  0;  0 |];
      [|-1;  0;  2; -1;  0;  0;  0 |];
      [| 0; -1; -1;  2; -1;  0;  0 |];
      [| 0;  0;  0; -1;  2; -1;  0 |];
      [| 0;  0;  0;  0; -1;  2; -1 |];
      [| 0;  0;  0;  0;  0; -1;  2 |]
    |]
  | E8 -> [|
      [| 2;  0; -1;  0;  0;  0;  0;  0 |];
      [| 0;  2;  0; -1;  0;  0;  0;  0 |];
      [|-1;  0;  2; -1;  0;  0;  0;  0 |];
      [| 0; -1; -1;  2; -1;  0;  0;  0 |];
      [| 0;  0;  0; -1;  2; -1;  0;  0 |];
      [| 0;  0;  0;  0; -1;  2; -1;  0 |];
      [| 0;  0;  0;  0;  0; -1;  2; -1 |];
      [| 0;  0;  0;  0;  0;  0; -1;  2 |]
    |]
  | F4 -> [|
      [| 2; -1;  0;  0 |];
      [|-1;  2; -2;  0 |];
      [| 0; -1;  2; -1 |];
      [| 0;  0; -1;  2 |]
    |]
  | G2 -> [|
      [| 2; -1 |];
      [|-3;  2 |]
    |]

let cartan_integer t i j = (cartan_matrix t).(i).(j)

let root_add_simple gamma j = gamma.(j) <- gamma.(j) + 1
let root_minus_simple gamma j = gamma.(j) <- gamma.(j) - 1
let root_equal = Array.equal Int.equal
let root_pairing t beta j =
  Array.fold_left Int.add 0
    (Array.mapi (fun i coeff -> coeff * cartan_integer t i j) beta)

let root_to_string beta =
  let first_term = ref true in
  let term_to_string i = function
    | 0 -> ""
    | 1 -> if !first_term then
             begin
             first_term := false;
             Printf.sprintf "𝛼%s" (subscript i)
             end
           else
             Printf.sprintf " + 𝛼%s" (subscript i)
    | c -> if !first_term then
             begin
             first_term := false;
             Printf.sprintf "%d 𝛼%s" c (subscript i)
             end
           else
             Printf.sprintf " + %d 𝛼%s" c (subscript i)
  in
  Array.fold_left String.cat "" (Array.mapi term_to_string beta)


let simple_root t i = Array.init (rank t) (function j -> if i = j then 1 else 0)
let simple_roots t = List.init (rank t) (simple_root t)

let rec root_mem beta = function
    [] -> false
  | alpha::l -> root_equal alpha beta || root_mem beta l

let rec nodup_append beta = function
    [] -> [beta]
  | alpha::l as ll ->
      if root_equal alpha beta then ll else alpha::(nodup_append beta l)

let nodup_append_ref beta lref = lref := nodup_append beta !lref

let roots_of_next_height t roots_of_less_height roots_of_this_height =
  let result = ref [] in
  let aux beta =
    for j = 0 to (rank t) - 1 do
      let r = ref 0 in
      let gamma = Array.copy beta in
      root_minus_simple gamma j;
      while root_mem gamma roots_of_less_height do
        incr r;
        root_minus_simple gamma j
      done;
      if !r > root_pairing t beta j then
        begin
        let gamma = Array.copy beta in
        root_add_simple gamma j;
        nodup_append_ref gamma result
        end
    done
  in
  List.iter aux roots_of_this_height;
  !result

let print_list_roots = List.iter (function beta -> Printf.printf "\t%s\n" (root_to_string beta))

let positive_roots t =
  let height = ref 0
  and previous_height = ref []
  and current_height = ref []
  and next_height = ref (simple_roots t) in
  Printf.printf "Enumerating the positive roots in the root system of type %s:\n" (name t);
  let next_count = ref (List.length !next_height) in
  while !next_count > 0 do
    incr height;
    Printf.printf "%d roots of height %d:\n" !next_count !height;
    print_list_roots !next_height;
    previous_height := List.append !previous_height !current_height;
    current_height := !next_height;
    next_height := roots_of_next_height t !previous_height !current_height;
    next_count := List.length !next_height
  done;
  let positive_roots = List.append !previous_height !current_height in
  Printf.printf "In total, %d positive roots.\n" (List.length positive_roots);
  positive_roots

(* Elements of the extended root poset come in three types:
  positive roots, negative roots and copies of simple roots *)
type ext_root = Pos of int array | Neg of int array | Dot of int


(* The partial order relation on the extended root poset *)
let rec ext_root_le t = function
  | Pos a -> begin function
      | Pos b -> Array.for_all2 (<=) a b
      | Dot _ -> false
      | Neg _ -> false
    end
  | Dot j -> begin function
      | Pos b -> ext_root_le t (Pos (simple_root t j)) (Pos b)
      | Dot k -> j = k
      | Neg b -> false
    end
  | Neg a -> begin function
      | Pos b -> begin
          let kref = ref 0 in
          while !kref < rank t && not (ext_root_le t (Neg a) (Dot !kref) && ext_root_le t (Dot !kref) (Pos b)) do
            incr kref
          done;
          not (!kref = rank t)
          end
      | Dot j -> begin
          let kref = ref 0 in
          while !kref < rank t && not (cartan_integer t j !kref <> 0 && ext_root_le t (Neg a) (Neg (simple_root t !kref))) do
            incr kref
          done;
          not (!kref = rank t)
          end
      | Neg b -> Array.for_all2 (<=) b a
    end

(* The underlying set of the extended root poset of type t *)

let matrix_ext_root_poset t =
  let pos_roots = positive_roots t in
(*  let pos_roots_rev = List.rev pos_roots in *)
  let elements = Array.concat [
(*    Array.map (function a -> Pos a) (Array.of_list pos_roots_rev); *)
    Array.map (function a -> Pos a) (Array.of_list pos_roots);
    Array.init (rank t) (function j -> Dot j);
    Array.map (function a -> Neg a) (Array.of_list pos_roots)]
  in
  let number = Array.length elements in
  let matrix = Array.init_matrix number number (fun i j -> ext_root_le t elements.(i) elements.(j)) in
  matrix

(*
let show_matrix =
  let aux v = Printf.printf (if v then "1 " else "0 ") in
  Array.iter (fun row ->
    Array.iter aux row;
    print_newline ())

let _ =
  let m = matrix_ext_root_poset E6 in
  show_matrix m
*)
